// Formal Verification Project
// Non-pipelined AES Unit

// Control Unit

// 11/18/2011



module control (clk, rst,
	 			roundkey_0, // input from key_expansion
				roundkey_1,
				roundkey_2,
				roundkey_3,
				roundkey_4,
				roundkey_5,
				roundkey_6,
				roundkey_7,
				roundkey_8,
				roundkey_9,
				roundkey_10,
				current_key, // output to add_roundkey containing the current round's key
				add_roundkey_o,  // input from add_roundkey
				sub_bytes_o, // input from sub_bytes
				shift_rows_o, // input from shift_rows
				mix_columns_o, // input from mix_columns
				state_out, // common output to each worker module
				data_in,  // input plaintext
				data_out, // completed cyphertext from npAES module
				done, 
				start); // indicates start of incoming plaintext
				

input clk, rst;
input start;
input [0:127] roundkey_0, roundkey_1, roundkey_2, roundkey_3, roundkey_4, roundkey_5, roundkey_6, roundkey_7, roundkey_8, roundkey_9, roundkey_10;
input [0:127] add_roundkey_o, sub_bytes_o, shift_rows_o, mix_columns_o, data_in;
output [0:127] state_out, current_key, data_out;
output reg done;

reg [0:127] t_data_out; 
reg [0:127] r_current_key;

// state parameters
parameter SUB_BYTES = 0, SHIFT_ROWS = 1, MIX_COLUMNS = 2, ADD_ROUNDKEY = 3, START = 4, FINISH = 5;

reg [3:0] cRnd; // stores the current round
reg [3:0] pstate, nstate; // stores the present & next states

// wires
wire [0:127] inputsel_o;
wire [0:127] pipeInput;
reg [1:0] sel;
reg [3:0] key_sel;
reg isFirstRound;
reg isLastRound;

 
assign pipeInput = (start == 1'b1) ? data_in : inputsel_o; 

pipereg _pipereg (.din(pipeInput), .dout(state_out), .clk(clk), .rst(rst));

inputselector _inputselector (	.add_roundkey_o(add_roundkey_o), 
								.sub_bytes_o(sub_bytes_o), 
								.shift_rows_o(shift_rows_o),
								.mix_columns_o(mix_columns_o),
								.state(inputsel_o),
								.sel(sel),
								.data_in(data_in));
								
keyselector _keyselector	(	.roundkey_0(roundkey_0),
								.roundkey_1(roundkey_1),
								.roundkey_2(roundkey_2),
								.roundkey_3(roundkey_3),
								.roundkey_4(roundkey_4),
								.roundkey_5(roundkey_5),
								.roundkey_6(roundkey_6),
								.roundkey_7(roundkey_7),
								.roundkey_8(roundkey_8),
								.roundkey_9(roundkey_9),
								.roundkey_10(roundkey_10),
								.current_key(current_key),
								.key_sel(key_sel));
								
																					

always @(posedge clk or posedge rst)
	begin
		if (rst == 1'b1)
			begin
				nstate <= START;
			end
		else
			begin
				pstate = nstate;
			end
	end // end always block
	
	
always @(posedge clk)
	begin
		case(pstate)
			START:
				begin
					done <= 0;
					if (start == 1'b0)
						begin
							nstate <= START; // keep waiting
						end
					else if (start == 1'b1)// start round 1
						begin
							cRnd <= 1;
							key_sel <= 0;
							nstate <= ADD_ROUNDKEY;
							isFirstRound <= 1;
							isLastRound <= 0;
						end
				end
			SUB_BYTES:
				begin
					if (isFirstRound == 0)
						begin
						cRnd <= cRnd + 1;
						end
					else if (isFirstRound == 1)
							begin
								isFirstRound = 0;
							end
					sel <= 2'b01;
					nstate <= SHIFT_ROWS;
				end
			SHIFT_ROWS:
				begin							
					sel <= 2'b10;
					if (cRnd == 4'b1010) // last round
						begin
							nstate <= ADD_ROUNDKEY;
							isLastRound <= 1;
						end
					else 
						begin
							nstate <= MIX_COLUMNS;
						end
				end
			MIX_COLUMNS:
				begin
					sel <= 2'b11;
					nstate <= ADD_ROUNDKEY;
				end
			ADD_ROUNDKEY:
				begin
					// output the roundkey by selecting the appropriate key_sel value
					sel <= 2'b00;
					if (isFirstRound == 1)
						begin
							nstate <= SUB_BYTES;
							key_sel <= 0;
							/*isFirstRound <= 0;*/
						end
					else if (isLastRound == 1)
						begin
							nstate <= FINISH;
							key_sel <= cRnd;
							isLastRound <= 0;
						end
					else if (isFirstRound == 0)
						begin
							nstate <= SUB_BYTES;
							key_sel <= cRnd;
							/*cRnd <= cRnd + 1;*/
						end
					end
			FINISH:
				begin
					// encryption is complete
					/*t_data_out <= state_out;*/ // output the cyphertext
					done <= 1; // indicate we are done
					nstate <= START; // go back the start state
				end
			default:
				begin
					
				end

		endcase
	
	
	
	end // end always block

	assign data_out = (done == 1'b1) ? state_out : 128'b0;

endmodule
	





















				 